import asyncio
import json
import os
import subprocess
import threading
import queue
from pathlib import Path
from typing import Dict, List, Optional

import aiofiles
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles


class ScriptManager:
    def __init__(self):
        self.active_processes: Dict[str, subprocess.Popen] = {}
        self.websocket_connections: Dict[str, WebSocket] = {}
        self.message_queues: Dict[str, queue.Queue] = {}
    
    def get_scripts_list(self) -> List[Dict]:
        """获取可用脚本列表"""
        scripts_dir = Path("../scripts")
        scripts = []
        if scripts_dir.exists():
            for script_file in scripts_dir.glob("*.py"):
                scripts.append({
                    "name": script_file.stem,
                    "filename": script_file.name,
                    "path": str(script_file)
                })
        return scripts
    
    async def start_script(self, script_name: str, websocket: WebSocket, session_id: str):
        """启动脚本执行"""
        script_path = Path(f"../scripts/{script_name}.py")
        if not script_path.exists():
            await websocket.send_text(json.dumps({
                "type": "error",
                "message": f"脚本 {script_name} 不存在"
            }))
            return
        
        try:
            print(f"Starting script: {script_name} for session {session_id}")
            
            # 设置环境变量以支持UTF-8编码
            env = os.environ.copy()
            env["PYTHONIOENCODING"] = "utf-8"
            env["PYTHONUNBUFFERED"] = "1"  # 禁用Python输出缓冲
            
            # 启动子进程
            process = subprocess.Popen(
                ["python", "-u", str(script_path)],  # -u 参数强制无缓冲
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                bufsize=0,  # 无缓冲
                universal_newlines=True,
                encoding='utf-8',
                cwd=script_path.parent,
                env=env
            )
            
            self.active_processes[session_id] = process
            self.websocket_connections[session_id] = websocket
            self.message_queues[session_id] = queue.Queue()
            
            # 创建输出读取线程
            def read_output():
                try:
                    import time
                    
                    while process.poll() is None:
                        # 在Windows上使用简单的readline方法
                        try:
                            line = process.stdout.readline()
                            if line:
                                self.message_queues[session_id].put({
                                    "type": "output",
                                    "data": line
                                })
                        except:
                            # 如果读取失败，稍微等待一下再继续
                            time.sleep(0.01)
                    
                    # 进程结束后读取剩余输出
                    try:
                        remaining_output = process.stdout.read()
                        if remaining_output:
                            self.message_queues[session_id].put({
                                "type": "output",
                                "data": remaining_output
                            })
                    except:
                        pass
                    
                    # 进程结束后发送结束信号
                    return_code = process.wait()
                    self.message_queues[session_id].put({
                        "type": "finished",
                        "return_code": return_code
                    })
                    
                except Exception as e:
                    print(f"Error reading output: {e}")
                    if session_id in self.message_queues:
                        self.message_queues[session_id].put({
                            "type": "error",
                            "message": f"读取输出错误: {str(e)}"
                        })
            
            # 启动输出读取线程
            threading.Thread(target=read_output, daemon=True).start()
            
            # 启动消息处理任务
            asyncio.create_task(self._process_messages(session_id))
            
            await websocket.send_text(json.dumps({
                "type": "started",
                "message": f"脚本 {script_name} 已启动"
            }))
            
        except Exception as e:
            await websocket.send_text(json.dumps({
                "type": "error",
                "message": f"启动脚本失败: {str(e)}"
            }))
    
    async def _process_messages(self, session_id: str):
        """处理来自子线程的消息"""
        while session_id in self.websocket_connections:
            try:
                # 非阻塞获取消息
                message = self.message_queues[session_id].get_nowait()
                websocket = self.websocket_connections[session_id]
                
                await websocket.send_text(json.dumps(message))
                
                # 如果是结束消息，清理资源
                if message.get("type") == "finished":
                    self._cleanup_session(session_id)
                    break
                    
            except queue.Empty:
                # 队列为空，等待一小段时间后继续
                await asyncio.sleep(0.01)
            except Exception as e:
                print(f"Error processing message: {e}")
                break
    
    def _cleanup_session(self, session_id: str):
        """清理会话资源"""
        if session_id in self.active_processes:
            del self.active_processes[session_id]
        if session_id in self.websocket_connections:
            del self.websocket_connections[session_id]
        if session_id in self.message_queues:
            del self.message_queues[session_id]
    
    async def send_input(self, session_id: str, user_input: str):
        """向脚本发送用户输入"""
        if session_id in self.active_processes:
            process = self.active_processes[session_id]
            try:
                process.stdin.write(user_input + "\n")
                process.stdin.flush()
            except Exception as e:
                websocket = self.websocket_connections.get(session_id)
                if websocket:
                    await websocket.send_text(json.dumps({
                        "type": "error",
                        "message": f"发送输入失败: {str(e)}"
                    }))
    
    def terminate_script(self, session_id: str):
        """终止脚本执行"""
        if session_id in self.active_processes:
            process = self.active_processes[session_id]
            process.terminate()
            try:
                process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                process.kill()
            
            self._cleanup_session(session_id)


app = FastAPI(title="脚本管理系统")

# 添加CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 静态文件服务
app.mount("/static", StaticFiles(directory="../frontend"), name="static")

script_manager = ScriptManager()


@app.get("/")
async def serve_frontend():
    """提供前端页面"""
    return FileResponse("../frontend/index.html")


@app.get("/api/scripts")
async def get_scripts():
    """获取可用脚本列表"""
    return {"scripts": script_manager.get_scripts_list()}


@app.get("/api/script/{script_name}")
async def get_script_content(script_name: str):
    """获取脚本内容"""
    script_path = Path(f"../scripts/{script_name}.py")
    if not script_path.exists():
        raise HTTPException(status_code=404, detail="脚本不存在")
    
    async with aiofiles.open(script_path, 'r', encoding='utf-8') as f:
        content = await f.read()
    
    return {"content": content}


@app.websocket("/ws/{session_id}")
async def websocket_endpoint(websocket: WebSocket, session_id: str):
    """WebSocket连接处理"""
    await websocket.accept()
    
    try:
        while True:
            data = await websocket.receive_text()
            message = json.loads(data)
            
            if message["type"] == "start_script":
                print(f"Starting script: {message['script_name']} for session {session_id}")
                script_name = message["script_name"]
                await script_manager.start_script(script_name, websocket, session_id)
            
            elif message["type"] == "input":
                user_input = message["data"]
                await script_manager.send_input(session_id, user_input)
            
            elif message["type"] == "terminate":
                script_manager.terminate_script(session_id)
                await websocket.send_text(json.dumps({
                    "type": "terminated",
                    "message": "脚本已终止"
                }))
                
    except WebSocketDisconnect:
        # 连接断开时清理资源
        script_manager.terminate_script(session_id)
    except Exception as e:
        print(f"WebSocket error: {e}")
        script_manager.terminate_script(session_id)


if __name__ == "__main__":
    import uvicorn
    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)
